#!/usr/bin/env python
#-*- coding: utf-8 -*-

# Author: Edward Roualdes
#         [2013.02.05]
#         University of Kentucky

from Bio import SeqIO
from os.path import splitext
import argparse
import sys

# Script to truncate reads, from the end, down to a specified length.
# $ python ./truncate.py --help 
# will show all options; all options are required


if __name__ == "__main__":

    # parse command line options
    parser = argparse.ArgumentParser(description='Truncate reads to length l by trimming off of the end of the sequences the appropriate number of nucleotides.')
    parser.add_argument("-i", "--input", 
                        help='/path/to/input/file', required=True)
    parser.add_argument("-o", "--output", help='/path/to/output/file')
    parser.add_argument("-l", "--length", 
                        type=int, help='int to truncate to', required=True)
    args = parser.parse_args()

    # if output option not specified
    # output file = input file with '_trun' concatenated and same extension
    if args.output:
        out = args.output
    else:
        out = splitext(args.input)
        out = out[0] + '_trun' + out[1]

    l = args.length             # specified length to truncate to
    trun = []

    # open input file
    fOpen = open(args.input, 'rU')
    for j in SeqIO.parse(fOpen, "fasta"): # parse read by read
        if len(j) > l:                    # truncate if needed
            trun.append(j[0:l]) 
        else:
            trun.append(j)
    fOpen.close()
    
    # write files to output
    output_handle = open(out, "w")
    SeqIO.write(trun, output_handle, "fasta")
    output_handle.close()
    
    sys.exit()
